# 雷达设计v1版智能体
from typing import Dict, List
import json
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.cpm_engine import CpmEngine
from core.radar_core import RadarCore
from common.app_manager import AppManager

class RadarDesignV1Agent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.pa_snr_agent.PaSnrAgent'

    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        observes = RadarDesignV1Agent._observe(context=context)
        plans = RadarDesignV1Agent._reason(observes=observes, context=context)
        return RadarDesignV1Agent._act(observes=observes, context=context, plans=plans)
    
    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        prompt = RadarDesignV1Agent._gen_params_parse_prompt(user_query=context.current_input)
        resp = CpmEngine.infer(query=prompt)
        print(f'resp json: \n{resp};')
        if "'''json\n" in resp:
            params = json.loads(resp[8:-3])
        else:
            params = json.loads(resp)
        context.states['params'].update(params)
        observes.update(params)
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = []
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
        # 生成试验报告
         # 创建目录: work/reports/v001
        base_folder = 'work/reports/e006'
        Path(base_folder).mkdir(parents=True, exist_ok=True)
        # 创建图片目录：work/reports/v001/images
        images_folder = f'{base_folder}/images'
        Path(images_folder).mkdir(parents=True, exist_ok=True)
        # 执行业务逻辑
        c = 3e8 # 光速
        k = 1.38e-23 # constant
        sigmaa = observes['sigmaa']
        sigmam = observes['sigmam']
        Theta_A = observes['Theta_A']
        T_sc = observes['T_sc']
        delta_R = observes['delta_R']
        nf_dB = observes['nf_dB']
        loss_dB = observes['loss_dB']
        te = observes['te']
        R_min = observes['R_min']
        v_i = observes['v_i']
        v_a = observes['v_a']
        v_m = observes['v_m']
        e_a = observes['e_a']
        e_m = observes['e_m']
        dec_SNR_dB = observes['dec_SNR_dB']
        # 1. 求出脉冲宽度和带宽
        tau = 2.0*delta_R/c
        B = 1.0 / tau
        # 2. 确定频率freq、孔径尺寸A_e、单脉冲峰值功率pt
        A_e = 2.25 # 有效孔径 m^{2}
        freq = 3e9      # 雷达频率 (Hz)
        # 求天线增益
        lambda_ = c / freq
        G = 4*np.pi*A_e / (lambda_**2)
        G_dB = 10*np.log10(G)
        # 计算脉冲重复频率
        t_interceptor = R_min / v_i # 防空导弹飞行时间（水平速度）
        tilde_R_a = R_min + t_interceptor * v_a # 飞机可被拦截最小距离
        tilde_R_m = R_min + t_interceptor * v_m # 导弹可被拦截最小距离
        scale_a = 90.0/78.0 # 为安全对飞机留有余量的系数
        scale_m = 55.0/48.0 # 为安全对导弹留有余量的系数
        R_a = tilde_R_a * scale_a
        R_m = tilde_R_m * scale_m
        R_u = max(R_a, R_m) # 取距离中的较大值为非模糊距离
        f_r = c / (2*R_u) # 计算脉冲重复频率
        f_r = int(f_r)  # 取一整数近似值
        # 计算竖直扫描范围
        theta_1 = np.rad2deg(np.atan(e_m / R_m))
        theta_2 = np.rad2deg(np.atan(e_a / R_min))
        print(f'c1: {theta_1}; c2: {theta_2};')
        # 计算扫描空间
        Theta_E = theta_2 - theta_1
        Omega = Theta_A * Theta_E / 57.296**2
        # 确定竖直、水平方向波束宽度
        theta_e = theta_2 - theta_1 # 竖直方向用一个波束覆盖
        theta_a = np.rad2deg((4*np.pi)/(G*np.deg2rad(theta_e)))
        # 确定累积波束数
        dot_theta_scan = Theta_A / T_sc
        n_P = int((theta_a * f_r)/dot_theta_scan)
        # A.非相参累积
        # A.1. 求单脉冲信噪比
        SNR_NCI_dB = dec_SNR_dB
        SNR_NCI = 10**(SNR_NCI_dB / 10.0)
        SNR_1 = SNR_NCI / (2*n_P) + np.sqrt(SNR_NCI**2/(4*n_P**2) + SNR_NCI/n_P)
        SNR_1_dB = 10*np.log10(SNR_1)
        print(f'SNR_1: {SNR_1}; {SNR_1_dB}dB;')
        # 计算非相参累积损失
        L_NCI = (1+SNR_1)/SNR_1
        L_NCI_dB = 10*np.log10(L_NCI)
        # 计算导弹单个脉冲能量
        print(f'R_m: {R_m}; G:{G}; lambda_:{lambda_}; sigmam: {sigmam};')
        nf = 10**(nf_dB/10.0)
        loss = 10**(loss_dB/10.0)
        E_m = ( (4*np.pi)**3*k*te*nf*loss*R_m**4*SNR_1 )/(G**2*lambda_**2*sigmam)
        E_a = ( (4*np.pi)**3*k*te*nf*loss*R_a**4*SNR_1 )/(G**2*lambda_**2*sigmaa)
        E = max(E_m, E_a)
        print(f'E_m: {E_m}; {E_a}; => {E};')
        # 求峰值功率
        P_t = E / tau
        print(f'pt: {P_t};')
        # 计算单脉冲SNR和积分损耗
        SNR_1 = (10**1.3)/(2*n_P) + np.sqrt(((10**1.3)**2)/(4*n_P*n_P) + (10**1.3)/n_P)
        LNCI = 10 * np.log10((1 + SNR_1)/SNR_1)
        loss_total = loss + LNCI
        # 生成距离范围 (15km到100km, 1000个点)
        range_vals = np.linspace(15e3, 100e3, 1000)  # 单位: 米
        range_km = range_vals / 1000  # 转换为公里
        # 计算SNR - 导弹目标
        snrm_nci = RadarCore.radar_eq(n_P*P_t, freq, G_dB, sigmam, te, B, nf, loss_total, range_vals)
        snrm_single = RadarCore.radar_eq(P_t, freq, G_dB, sigmam, te, B, nf, loss, range_vals)
        # 计算SNR - 飞机目标
        snra_nci = RadarCore.radar_eq(n_P*P_t, freq, G_dB, sigmaa, te, B, nf, loss_total, range_vals)
        snra_single = RadarCore.radar_eq(P_t, freq, G_dB, sigmaa, te, B, nf, loss, range_vals)
        # 创建图形
        plt.figure(figsize=(10, 8))
        # 导弹情况 - 上子图
        plt.subplot(2, 1, 1)
        plt.plot(range_km, snrm_nci, 'k-', label='With 7-pulse NCI')
        plt.plot(range_km, snrm_single, 'k-.', label='Single pulse')
        plt.grid(True)
        plt.legend()
        plt.ylabel('SNR - dB')
        plt.title('Missile case')
        # 飞机情况 - 下子图
        plt.subplot(2, 1, 2)
        plt.plot(range_km, snra_nci, 'k-', label='With 7-pulse NCI')
        plt.plot(range_km, snra_single, 'k-.', label='Single pulse')
        plt.grid(True)
        plt.legend()
        plt.ylabel('SNR - dB')
        plt.xlabel('Detection range - Km')
        plt.title('Aircraft case')
        plt.tight_layout()
        # plt.show()
        plt.savefig(f'{images_folder}/f01.png')
        # 更新上下文对象
        context.states['params'].update(observes)
        # plt.savefig(f'{images_folder}/f00.png')
        # 生成报告文本
        with open(f'{base_folder}/report.md', 'w', encoding='utf-8') as wfd:
            wfd.write(f'# 雷达设计v1版\n')
            wfd.write(f'# 1. 问题描述\n')
            wfd.write(f'{context.current_input}\n')
            wfd.write(f'# 2. 计算过程\n')
            wfd.write(f'## 2.1. 求脉冲宽度和带宽\n')
            wfd.write(f'脉冲宽度：$\\tau={tau}$；带宽：B={B}\n')
            wfd.write(f'## 2.2. 天线增益\n')
            wfd.write(f'G={G_dB:0.2f}dB\n')
            wfd.write(f'## 2.3. 计算脉冲重复频率$f_{{r}}$\n')
            wfd.write(f'系统如下图所示：\n![f00.png](./images/f00.png "")\n')
            wfd.write(f'防空导弹飞行时间（仅考虑水平速度分量）：\n')
            wfd.write(f'$$\nt_{{interceptor}}=\\frac{{R_{{min}}}}{{v_{{i}}}}={t_interceptor}\n$$\n')
            wfd.write(f'飞机可被拦截最小距离：$\\tilde{{R_a}} = R_{{min}} + t_{{interceptor}} \\times v_{{a}} = {tilde_R_a}$\n')
            wfd.write(f'导弹可被拦截最小距离：$\\tilde{{R_m}} = R_{{min}} + t_{{interceptor}} \\times v_{{m}} = {tilde_R_m}$\n')
            wfd.write(f'设飞机的安全余量系数为{scale_a}，导弹的安全余量系数为{scale_m}。则飞机拦截距离为{R_a}，导弹拦截距离为{R_m}，取二者最大者为非模糊距离$R_u={R_u}$\n')
            wfd.write(f'脉冲重复频率：$f_r=\\frac{{c}}{{2 \\times R_u}} \\approx {f_r}$\n')
            wfd.write(f'## 2.4. 计算扫描空间\n')
            wfd.write(f'$\\Omega={Omega}$\n')
            wfd.write(f'## 2.5. 确定波束宽度\n')
            wfd.write(f'水平波束宽度：{theta_a}，竖直波束宽度：{theta_e}。\n')
            wfd.write(f'## 2.6. 累积波束数\n')
            wfd.write(f'$n_P={n_P}$\n')
            wfd.write(f'## 2.7. 单脉冲信噪比\n')
            wfd.write(f'SNR_1_dB={SNR_1_dB}dB\n')
            wfd.write(f'## 2.8. 非相参累积损失\n')
            wfd.write(f'L_NCI_dB={L_NCI_dB}dB')
            wfd.write(f'## 2.9 求峰值功率\n')
            wfd.write(f'$p_t={P_t}\n')
            wfd.write(f'## 2.10. 总损耗\n')
            wfd.write(f'计算得到的积分损耗: {LNCI:.2f} dB, 总系统损耗: {loss_total:.2f} dB\n')
            wfd.write(f'## 2.11. 单脉冲和非相参累积比较图\n')
            wfd.write(f'![f01.png](./images/f01.png "")\n')
        return {'result': 0, 'params': {}}
    
    @staticmethod
    def _gen_params_parse_prompt(user_query:str) -> str:
        template = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户输入，提取出所有的参数，并转换为标准单位，以json格式输出。由于参数很多，而且参数表述和单位形式多样，你需要按照下面的步骤执行。

# 要求
你只需要给出json结果，不需要生成代码或进行解释。

# 步骤
1. 飞机RCS（sigmaa）：又叫截面积，单位为$m^{2}$或dBsm，如果单位为dBsm需转为$m^{2}$，公式为：m=10**(sigmaa/10)，提取值为v1；
2. 导弹RCS（sigmam）：又叫截面积，单位为$m^{2}$或dBsm，如果单位为dBsm需转为$m^{2}$，公式为：m=10**(sigmaa/10)，提取值为v2；
3. 水平扫描范围（Theta_A）：单位为$\\circ$（度），提取值为v3；
4. 扫描时间（T_sc）：通过扫描速度和转一圈附近内容确定，单位为s（秒）、ms（毫秒）、us（微秒），若单位为ms时需乘以e-3，若单位为us时需乘以e-6，提取值为v4；
5. 距离分辨率（delta_R）：，单位为m（米），提取值为v5；
6. 噪声系数（nf_dB）：又叫噪声数，单位为dB（分贝），若仅为数值应转为dB，转换公式为：dB=10*np.log10(nf)，提取值为v6；
7. 系统损耗（loss_dB）：又叫总接收损失，单位为dB（分贝），若仅为数值应转为dB，转换公式为：dB=10*np.log10(nf)，提取值为v7；
8. 噪声温度（te）：又叫工作温度、环境温度，单位为K（开尔文），若单位为$\\circ$（度）及摄氏度时需加上275.12，提取值为v8；
9. 净空距离（R_min）：单位为m（米）、Km（千米、公里），若单位为Km时需乘以e3，提取值为v9；
10. 防空导弹速度（v_i）：又叫拦截导弹速度，单位：m/s（米/秒）、Km/h（千米/小时），当单位为Km/h时需v_i*e3/3600，提取值为v10；
11. 飞机速度（v_a）：单位：m/s（米/秒）、Km/h（千米/小时），当单位为Km/h时需v_a*e3/3600，提取值为v11；
12. 导弹速度（v_m）：单位：m/s（米/秒）、Km/h（千米/小时），当单位为Km/h时需v_m*e3/3600，提取值为v12；
13. 飞机飞行高度（e_a）：单位为m（米）、Km（千米、公里），若单位为Km时需乘以e3，提取值为v13；
14. 导弹飞行高度（e_m）：单位为m（米）、Km（千米、公里），若单位为Km时需乘以e3，提取值为v14；
15. 检测阈值信噪比（dec_SNR_dB）：提取方法检测阈值和信噪比联合出现，单位为dB（分贝），若仅为数值应转为dB，转换公式为：dB=10*np.log10(nf)，提取值为v15；
16. 生成json结果：
{
    "sigmaa": v1,
    "sigmam": v2,
    "Theta_A": v3,
    "T_sc": v4,
    "delta_R": v5,
    "nf_dB": v6,
    "loss_dB": v7,
    "te": v8,
    "R_min": v9,
    "v_i": v10,
    "v_a": v11,
    "v_m": v12,
    "e_a": v13,
    "e_m": v14,
    "dec_SNR_dB": v15
}

# 示例
输入：设计一款满足如下功能的搜索和检测地基雷达。威胁包括飞机和导弹两大类。其中飞机的RCS为6dBsm（$\\sigma_{a}=4m^{2}$），导弹平均RCS为-3dBsm（$\\sigma_{m}=0.5m^{2}$）。导弹的高度是2Km，飞机的高度为7Km。水平扫描范围为360度，扫描速度为2秒钟转一圈。采用L到X波段，距离分辨率为150m。角度分辨率暂时不要求。危胁由一个飞机和一个导弹构成。噪声数为$F=6dB$，总接收损失$L=8dB$。采用fan beam方式，水平扫描波束宽度小于3度。假设13dB的信噪比是检测阈值，净空距离为$R_{min}=30Km$，防空导弹的速度为$v_{i}=250m/s$，飞机的速度为$v_{a}=400m/s$，导弹速度为$v_{m}=150m/s$，并且地面是平的，环境温度为290K。
输出：
{
    "sigmaa": 4.0,
    "sigmam": 0.5,
    "Theta_A": 360.0,
    "T_sc": 2.0,
    "delta_R": 150,
    "nf_dB": 6.0,
    "loss_dB": 8.0,
    "te": 290.0,
    "R_min": 30e3,
    "v_i": 250.0,
    "v_a": 400.0,
    "v_m": 150.0,
    "e_a": 7e3,
    "e_m": 2e3,
    "dec_SNR_dB": 13.0
}


# 当前输入
'''
        prompt = f'{template}{user_query}'
        return prompt